import os

import numpy as np
import torch

from util.seed import set_seed
from raw.get_model import get_model
from attack.get_attack import get_attack

# 生成并保存攻击后图像
def main(model_name, attack_name, attack_config, suffix = ''):
    # 设置随机数种子
    set_seed()

    # 创建网络
    model = get_model(model_name)
    # 加载已保存的网络
    model.load_raw_state()

    # 加载攻击
    # attack_config = attack_config.copy()
    # attack_config['suffix'] = suffix
    attack = get_attack(attack_name, model, attack_config)

    # 加载选取的数据
    base_dir = 'data/candidate_data/' + model_name + '/'
    images = np.load(base_dir + 'images.npy')
    labels = np.load(base_dir + 'labels.npy')
    if attack.targeted:
        if attack.llc:
            targets = np.load(base_dir + 'llcs.npy')
        else:
            targets = np.load(base_dir + 'targets.npy')
    else:
        targets = labels

    #进行攻击，生成混淆图像
    adv_images = attack.attack(images, targets)

    #将图像传入网络获取标签
    adv_labels = np.argmax(model.predict(adv_images), axis = 1)

    #计算获取标签错误率
    mis = 0
    for i in range(len(adv_images)):
        if labels[i] != adv_labels[i]:
            mis = mis + 1
    print('Misclassification ratio is {}/{}={:.1f}%.'.format(
        mis, len(adv_images), mis / len(adv_images) * 100))

    # 创建保存目录
    name = model_name + '_' + attack_name + suffix
    outer_dir = 'data/adversarial_data/'
    base_dir = outer_dir + name + '/'
    if name not in os.listdir(outer_dir):
        os.mkdir(base_dir)

    #保存混淆图像与识别出的标签
    np.save(base_dir + 'images.npy', adv_images)
    np.save(base_dir + 'labels.npy', adv_labels)
    print('Adversarial data saved to ' + base_dir + '*.npy')
